Skip to content

Track LLM token usage in reports#244

Open
HexSleeves wants to merge 2 commits into
NVIDIA:mainfrom
HexSleeves:codex/add-llm-token-usage-reporting
Open

Track LLM token usage in reports#244
HexSleeves wants to merge 2 commits into
NVIDIA:mainfrom
HexSleeves:codex/add-llm-token-usage-reporting

Conversation

@HexSleeves

Copy link
Copy Markdown

Motivation

  • Capture per-call token telemetry so report metadata can surface LLM token consumption and detect cost/usage trends.
  • Preserve existing structured-output parsing while retaining raw provider metadata needed for usage (LangChain raw payload).

Description

  • Add token counters to LLMCallRecord and llm_call_record() defaults: input_tokens, output_tokens, and total_tokens.
  • Use with_structured_output(..., include_raw=True) and unwrap {"raw","parsed","parsing_error"} responses, extracting usage from raw.usage_metadata (supports input_tokens/output_tokens and prompt_tokens/completion_tokens variants) and preserving parsing-error behavior.
  • Record token usage in both sync (run_batches) and async (arun_batches) paths (including non-structured raw responses), and propagate analyzer.llm_usage into node-level llm_call_log entries (success and failure).
  • Aggregate llm_call_log token counters into report JSON metadata under metadata.llm_usage while keeping existing metadata intact.
  • Add support in the CLI structured-output adapter for include_raw=True so local/CLI providers expose raw usage without changing parser expectations.
  • Add unit tests covering the new fields, sync/async usage capture, missing usage metadata handling, structured-output include_raw=True parsing, and metadata aggregation; closes feat: expose LLM token usage in JSON report output #242.

Testing

  • Ran ruff format . and ruff check . with no issues.
  • Ran the full test suite with pytest -q: all tests passed (1261 passed, 12 skipped, 34 deselected, 6 xfailed).
  • mypy src was executed and reported pre-existing typing issues in unrelated modules; no new type regressions introduced by these changes (mypy warnings are repository pre-existing and not caused by this PR).

@HexSleeves HexSleeves force-pushed the codex/add-llm-token-usage-reporting branch from 3bb9d24 to 622ab7d Compare July 6, 2026 02:16

@rng1995 rng1995 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Automated SkillSpector Review]

Requesting changes. Usage capture works for LLMAnalyzerBase, but the report still undercounts the real TP4 graph LLM call. Instrument the direct chat_completion() path and add a regression asserting provider usage is included in metadata.llm_usage.

Comment thread tests/test_mcp_tool_poisoning.py Outdated
):
result = node(state)
assert result["llm_call_log"] == [{"node": "mcp_tool_poisoning", "ok": True, "error": None}]
assert result["llm_call_log"] == [llm_call_record("mcp_tool_poisoning", ok=True)]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Blocking: this expectation bakes in zero tokens for TP4 even though production mcp_tool_poisoning.py calls chat_completion() directly and records a successful LLM call. That path bypasses the new raw-response usage extraction, so report totals are incomplete. Instrument chat_completion/TP4 and assert the provider's nonzero usage metadata here.

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 8b45e81. chat_completion_with_usage() now preserves normalized provider token metadata, TP4 forwards it into llm_call_log (including parse failures), and the regression verifies the nonzero TP4 usage reaches metadata.llm_usage. Verification: Ruff clean; 1262 tests passed.

@rng1995 rng1995 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Automated SkillSpector Review]

Re-review of head 8b45e81. The PR adds token counters to LLMCallRecord, captures provider usage in LLMAnalyzerBase (sync/async, structured via include_raw=True and raw paths), instruments the TP4 chat_completion path via a new chat_completion_with_usage(), and aggregates totals into report JSON under metadata.llm_usage.

Prior-issue resolution checklist

  • TP4 chat_completion() path bypassed usage extraction; test baked in zero tokens → Resolved. Commit 8b45e81 adds chat_completion_with_usage() in src/skillspector/llm_utils.py returning (content, usage); _check_tp4 now records **usage on both the ok record and the post-attempt failure record (so usage from a successful call whose JSON fails to parse is still counted). The regression test_successful_call_records_provider_usage asserts nonzero provider usage (12/3/15) lands in the llm_call_log record and flows through _build_metadata into metadata.llm_usage. chat_completion() is kept as a thin back-compat wrapper, and the integration test's patch target was updated.

New-commit scan (8b45e81)

No new blockers. The token-usage helpers were sensibly moved from llm_analyzer_base into llm_utils (LLMTokenUsage, empty_token_usage, extract_token_usage) so both transports share one normalizer. Parse-failure behavior is preserved (_unwrap_structured_response re-raises parsing_error, so the ValidationError degradation branches still trigger), and the CLI adapter's fail-closed ValueError on garbage JSON is unchanged for the default include_raw=False path.

Schema contract

  • metadata.llm_usage is additive in the JSON report; terminal/markdown/SARIF outputs are untouched.
  • LLMCallRecord gains three int fields, but every producer goes through llm_call_record() (defaults 0) and test assertions were migrated to the builder, so the log shape stays consistent. Non-breaking.

Non-blocking nits

  • extract_token_usage (src/skillspector/llm_utils.py:73-75): int(...) on malformed provider metadata can raise inside the live call path — see inline comment.
  • The _NoUsage shim + locals().get("analyzer", _NoUsage()) pattern is duplicated in four modules (semantic_developer_intent.py, semantic_quality_policy.py, semantic_security_discovery.py, meta_analyzer.py). Initializing analyzer: LLMAnalyzerBase | None = None before the try and using analyzer.llm_usage if analyzer else empty_token_usage() would be clearer and avoid the shared mutable class-attribute dict.
  • The PR description claims a test for "structured-output include_raw=True parsing", but _StructuredAgentCLIModel.invoke with include_raw=True (the dict-returning branch, including a parse failure captured as parsing_error) has no direct test — worth adding alongside test_structured_output_parses_and_validates in tests/unit/test_llm_utils.py.

usage = getattr(raw, "usage_metadata", None) or {}
if not isinstance(usage, dict):
return empty_token_usage()
input_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Non-blocking: these int(...) conversions run inline in the live LLM call path (now including TP4 via chat_completion_with_usage). If a provider ever returns a non-numeric usage value (e.g. "input_tokens": "n/a"), the raised ValueError/TypeError would convert a successful analyzer call into a failed one — degrading the scan and dropping findings over telemetry. Consider wrapping the conversions in a try/except (TypeError, ValueError): return empty_token_usage() so usage extraction can never fail the detection path.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

feat: expose LLM token usage in JSON report output

2 participants